from migflow import fluid
from migflow import scontact

import numpy as np
import os
import sys
import time
import gmsh
import shutil

outputdir = "output" if len(sys.argv) < 2 else sys.argv[1]
clscale = float(sys.argv[2]) if len(sys.argv) > 2 else 15

L = 0.04
H = 0.4
r = 5e-3 / 2
rho = 996
nu = 8e-7
mu = rho * nu
rhop = 1.01 * rho
xp = np.array((0, H / 2 - 4 * r))
lc = L * 1e-3 * clscale
lc_max = lc * 10
lc_min = lc

Re = 156
vRef = Re * mu / (rho * 2 * r)
vRef = 0.02501

print("vRef = %g" % vRef)


def size_callback(x, y):
    # return lc
    d = np.abs(x - xp[0])
    d = d / L

    distmin = 75 * r
    alpha = np.clip((d - distmin) / (10 * distmin), 0, 1)
    size = lc_min * (1 - alpha) + lc_max * alpha
    return size
    # linear
    return np.maximum(lc_max * np.maximum(d - 0.001, 0.0), lc_min)
    if d < 0.1:
        return lc_min
    return lc_max


def gen_mesh(h, w, mesh_size, origin=np.array([0, 0])):
    origin = np.asarray(origin)
    gmsh.model.add("box")

    p1 = gmsh.model.geo.add_point(-w / 2, +h / 2, 0, mesh_size)
    p2 = gmsh.model.geo.add_point(-w / 2, -h / 2, 0, mesh_size)
    p3 = gmsh.model.geo.add_point(+w / 2, -h / 2, 0, mesh_size)
    p4 = gmsh.model.geo.add_point(+w / 2, +h / 2, 0, mesh_size)
    l1 = gmsh.model.geo.add_line(p1, p2)
    l2 = gmsh.model.geo.add_line(p2, p3)
    l3 = gmsh.model.geo.add_line(p3, p4)
    l4 = gmsh.model.geo.add_line(p4, p1)
    line_loop = gmsh.model.geo.add_curve_loop([l1, l2, l3, l4])
    surface = gmsh.model.geo.add_plane_surface([line_loop])
    gmsh.model.geo.synchronize()
    gmsh.model.add_physical_group(1, [2], name="Bottom")
    gmsh.model.add_physical_group(1, [4], name="Top")
    gmsh.model.add_physical_group(1, [1, 3], name="Lateral")
    gmsh.model.add_physical_group(2, [1], name="domain")
    gmsh.model.mesh.set_size_callback(lambda dim, tag, x, y, z, lc: size_callback(x, y))
    gmsh.option.set_number("Mesh.Algorithm", 1)
    gmsh.model.mesh.generate(2)


gen_mesh(H, L, lc)

t = 0
ii = 0

print("Start with r=%g" % (r / 2))
shutil.rmtree(outputdir, True)
if not os.path.isdir(outputdir):
    os.makedirs(outputdir)
# physical parameters
g = np.array([0, -9.81])  # gravity
tEnd = 60 * 2 * r / vRef  # final time
tEnd_adim = tEnd * vRef / (2 * r)

# numerical parameters
cfl = 0.001 / 10
dt = cfl / vRef  # time step
outf = 10

p = scontact.ParticleProblem(2)
p.add_particle(xp, r, r**2 * np.pi * rhop)
p.write_mig(outputdir, 0)

f = fluid.FluidProblem(2, g, nu * rho, rho)
f.load_msh(None)
f.set_symmetry_boundary("Bottom")
f.set_symmetry_boundary("Lateral")
f.set_symmetry_boundary("Top")
f.set_mean_pressure(0)
f.set_particles(
    p.delassus(), p.volume(), p.position(), p.velocity(), p.omega(), p.contact_forces()
)
f.write_mig(outputdir, 0)
# Initial time and iteration
t = 0
ii = 0
# COMPUTATION LOOP
tic = time.perf_counter()
mass = p.r() ** 2 * np.pi * rhop
dp = []
drag = []
total_x = []
total_y = []
pvelocity_x = []
pvelocity_y = []
while t < tEnd:
    t += dt
    t_adim = t * vRef / (2 * r)
    print("t = %.2g" % t)
    print("t_adim = %.2g" % t_adim)
    f.set_particles(
        p.delassus(),
        p.volume(),
        p.position() + p.velocity() * dt,
        p.velocity(),
        p.omega(),
        p.contact_forces(),
    )
    f.implicit_euler(dt)
    pforce = f.compute_node_force()  # -drag-dp
    if ii % outf == 0:
        f.write_mig(outputdir, t_adim)
    p.iterate(dt, pforce + g * mass)
    # Output files writting
    if ii % outf == 0:
        p.write_mig(outputdir, t_adim)
    ii += 1
    dp_i = np.array([0.0])
    dp.append(dp_i)
    drag.append(-(pforce[:, 1] + dp_i))
    total_y.append(pforce[:, 1])
    total_x.append(pforce[:, 0])
    pvelocity_y.append(p.velocity()[-1, 1])
    pvelocity_x.append(p.velocity()[-1, 0])
last_velocity = p.velocity()[-1, 1]
print(last_velocity)

dp = np.array(dp).reshape(-1)
drag = np.array(drag).reshape(-1)
weight = g[1] * np.pi * r**2 * rhop * np.ones(dp.shape[0])
total_y = np.array(total_y).reshape(-1)
total_x = np.array(total_x).reshape(-1)

i = np.arange(drag.shape[0])
fmax = np.max(np.abs(total_y))
drag /= fmax
dp /= fmax
total_y /= fmax
total_x /= fmax
weight /= fmax
pvelocity_y = np.abs(np.array(pvelocity_y).reshape(-1)) / vRef
pvelocity_x = np.array(pvelocity_x).reshape(-1) / vRef


simu_time = np.arange(dp.shape[0]) * dt
np.savetxt("time.txt", simu_time)
np.savetxt("drag.txt", drag)
np.savetxt("dp.txt", dp)
np.savetxt("weight.txt", weight)
np.savetxt("total_y.txt", total_y)
np.savetxt("total_x.txt", total_x)
np.savetxt("pvelocity_y.txt", pvelocity_y)
np.savetxt("pvelocity_x.txt", pvelocity_x)

PLOTTING = False
if PLOTTING:
    import matplotlib.pyplot as plt

    fig, ax = plt.subplots(1, 2)
    buoyancy = (rhop - rho) * g[1] * np.ones(dp.shape[0]) * np.pi * r**2 / fmax
    ax[0].plot(i, -dp, "red", label="dp")
    ax[0].plot(i, -drag, "blue", label="drag")
    ax[0].plot(i, weight, "green", label="weight")
    ax[0].plot(i, total_y, "--k", label="total")
    ax[0].plot(i, total_x, "--", color="red", label="total")
    ax[0].set_ylim([-2, 2])
    ax[0].legend()

    ax[1].plot(i, pvelocity_y)
    ax[1].plot(i, pvelocity_x)
    ax[1].set_ylim([-0.2, 1.2])
    plt.show()